from os import path
import numpy as np
import pandas as pd
import torch
import torch.optim as optim
import matplotlib.pyplot as plot

DATA_PATH = path.realpath('pytorch/jizhi/bike/data/hour.csv')

class Bike():
    def exec(self):
        # self.prepare_data_and_params()
        self.prepare_data_and_params_2()
        # self.train_and_plot()
        self.train_and_plot2()
        # self.predict_and_plot()
        self.predict_and_plot2()
        
    def prepare_data_and_params(self):
        self.data = pd.read_csv(DATA_PATH)
        
        counts = self.data['cnt'][:50]
        self.x = torch.FloatTensor(np.arange(len(counts), dtype=float) / len(counts))
        self.y = torch.FloatTensor(np.array(counts, dtype=float))
        self.size = 10
        self.weights = torch.randn((1, self.size), requires_grad=True)
        self.biases = torch.randn((self.size), requires_grad=True)
        self.weights2 = torch.randn((self.size, 1), requires_grad=True)
        
    def prepare_data_and_params_2(self):
        # type columns to dummy
        self.data = pd.read_csv(DATA_PATH)
        self.rides = pd.read_csv(DATA_PATH)
        dummy_fields = ['season', 'weathersit', 'mnth', 'hr', 'weekday']
        for each in dummy_fields:
            dummies = pd.get_dummies(self.data[each], prefix=each, drop_first=False)
            self.data = pd.concat([self.data, dummies], axis=1)
        drop_fields = ['season', 'weathersit', 'mnth', 'hr', 'weekday', 'instant', 'dteday', 'workingday', 'atemp']
        self.data = self.data.drop(drop_fields, axis=1)
        
        # decimal columns to scaler
        quant_features = ['cnt', 'temp', 'hum', 'windspeed']
        self.scaled_features = {}
        for each in quant_features:
            mean, std = self.data[each].mean(), self.data[each].std()
            self.scaled_features[each] = [mean, std]
            self.data.loc[:, each] = (self.data[each] - mean) / std
            
        self.tr, self.te = self.data[:-21 * 24], self.data[-21 * 24:]
        target_fields = ['cnt', 'casual', 'registered']
        self.xtr, self.ytr = self.tr.drop(target_fields, axis=1), self.tr[target_fields]
        self.xte, self.yte = self.te.drop(target_fields, axis=1), self.te[target_fields]
        self.x = self.xtr.values.astype(float)
        y = self.ytr['cnt'].values.astype(float)
        self.y = np.reshape(y, [len(y), 1])        
        self.loss = []

    def train_and_plot(self):
        rate = 0.001
        losses = []
        x, y = self.x.view(50, -1), self.y.view(50, -1) # reshape
        for num in range(100000):
            hidden = x * self.weights + self.biases
            hidden = torch.sigmoid(hidden)
            predictions = hidden.mm(self.weights2)
            loss = torch.mean((predictions - y) ** 2)
            losses.append(loss.data.numpy())
            if num % 10000 == 0:
                print(f'loss: {loss}')
            loss.backward()
            
            self.weights.data.add_(- rate * self.weights.grad.data)
            self.biases.data.add_(- rate * self.biases.grad.data)
            self.weights2.data.add_(- rate * self.weights2.grad.data)
            
            self.weights.grad.data.zero_()
            self.biases.grad.data.zero_()
            self.weights2.grad.data.zero_()
        
        # plot loss
        '''
        plot.plot(losses)
        plot.xlabel('epoch')
        plot.ylabel('loss')
        plot.show()
        '''
        
        # plot predict
        '''
        x_data = x.data.numpy()
        plot.figure(figsize=(10, 7))
        xplot, = plot.plot(x_data, y.data.numpy(), 'o')
        yplot, = plot.plot(x_data, predictions.data.numpy())
        plot.xlabel('x')
        plot.ylabel('y')
        plot.legend([xplot, yplot], ['Data', 'prediction with 30000 epoch'])
        plot.show()
        '''
        
    def train_and_plot2(self):
        input_size = self.xtr.shape[1]
        hidden_size=10
        output_size=1
        batch_size=128
        self.neu = torch.nn.Sequential(
            torch.nn.Linear(input_size, hidden_size),
            torch.nn.Sigmoid(),
            torch.nn.Linear(hidden_size, output_size)
        )
        cost = torch.nn.MSELoss()
        optimizer = torch.optim.SGD(self.neu.parameters(), lr=0.01)
        
        for i in range(400):
            batch_loss = []
            for start in range(0, len(self.x), batch_size):
                end = start + batch_size if start + batch_size < len(self.x) else len(self.x)
                xx = torch.FloatTensor(self.x[start:end])
                yy = torch.FloatTensor(self.y[start:end])
                predictions = self.neu(xx)
                loss = cost(predictions, yy)
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
                batch_loss.append(loss.data.numpy())
            if i % 100 == 0:
                self.loss.append(np.mean(batch_loss))
                print(i, np.mean(batch_loss))
        # plot.plot(np.arange(len(self.loss)) * 100, self.loss)
        # plot.xlabel('epoch')
        # plot.ylabel('MSE')
        # plot.show()
    
    def predict_and_plot(self):
        counts_predict = self.data['cnt'][50:100]
        x = torch.FloatTensor((np.arange(len(counts_predict), dtype=float) + 50) / 100)
        y = torch.FloatTensor(np.array(counts_predict, dtype=float))
        
        # num multiply replace matrix multiply
        hidden = x.expand(self.size, len(x)).t() * self.weights.expand(len(x), self.size)
        hidden = torch.sigmoid(hidden)
        predictions = hidden.mm(self.weights2)
        loss = torch.mean((predictions - y) ** 2)
        print(f'loss: {loss}')
        
        x_data = x.data.numpy()
        plot.figure(figsize=(10, 7))
        xplot, = plot.plot(x_data, y.data.numpy(), 'o')
        yplot, = plot.plot(x_data, predictions.data.numpy())
        plot.xlabel('x')
        plot.ylabel('y')
        plot.legend([xplot, yplot], ['data', 'prediction'])
        plot.show()
        
    def predict_and_plot2(self):
        targets = self.yte['cnt']
        targets = targets.values.reshape([len(targets), 1]).astype(float)
        x = torch.FloatTensor(self.xte.values.astype(float))
        y = torch.FloatTensor(targets)
        predict = self.neu(x)
        predict = predict.data.numpy()
        
        fig, ax = plot.subplots(figsize=(10, 7))
        mean, std = self.scaled_features['cnt']
        ax.plot(predict * std + mean, label='prediction')
        ax.plot(targets * std + mean, label='data')
        ax.legend()
        ax.set_xlabel('date-time')
        ax.set_ylabel('counts')
        dates = pd.to_datetime(self.rides.loc[self.te.index]['dteday'])
        dates = dates.apply(lambda d: d.strftime('%b %d'))
        ax.set_xticks(np.arange(len(dates))[12::24])
        ax.set_xticklabels(dates[12::24], rotation=45)
        plot.show()

def main():
    Bike().exec()

if __name__ == '__main__':
    main()